#ifdef __cplusplus
extern "C" {
#endif // __cpusplus

#include <ntifs.h>

#ifdef __cplusplus
}
#endif // __cpusplus

#include "code_protection.h"
#include "dasm.h"
#define PROTECTED_DATA_TAG 'tdtp'

typedef struct _PROTECTED_DATA 
{
    LIST_ENTRY list_entry;
    PUCHAR data_address;
    ULONG data_size;
    UCHAR mirror_data[1];
}PROTECTED_DATA, *PPROTECTED_DATA;

static HANDLE m_thread_guard = NULL;
static HANDLE m_thread_back_guard = NULL;
static LIST_ENTRY m_protected_data_list;
static BOOLEAN m_is_protected = TRUE;
static ULONG m_meltdown_count = 0;

static VOID _guard_thread(
    IN PVOID StartContext
    );

static VOID _back_guard_thread(
    IN PVOID StartContext
    );

NTSTATUS start_protection()
{
    NTSTATUS status = STATUS_SUCCESS;
    OBJECT_ATTRIBUTES thread_attribute = {0};
    PUCHAR thread_function = NULL;
    ULONG code_size = 0;

    return STATUS_SUCCESS;

    InitializeListHead(&m_protected_data_list);

    InitializeObjectAttributes(
        &thread_attribute, 
        NULL, 
        OBJ_KERNEL_HANDLE, 
        NULL, 
        NULL);

    status = PsCreateSystemThread(
        &m_thread_guard,
        THREAD_ALL_ACCESS,
        &thread_attribute,
        NULL,
        NULL,
        _guard_thread,
        NULL);

    status = PsCreateSystemThread(
        &m_thread_back_guard,
        THREAD_ALL_ACCESS,
        &thread_attribute,
        NULL,
        NULL,
        _back_guard_thread,
        NULL);
    code_size = 0;
    thread_function = (PUCHAR)_back_guard_thread;
    while (thread_function[code_size] == 0xcc)
    {
        code_size += GetOpCodeSize(thread_function + code_size);
    }
    ++ code_size;
    add_protection(thread_function, code_size);

    return STATUS_SUCCESS;
}

VOID add_protection( 
    IN PUCHAR DataAddress, 
    IN ULONG DataSize 
    )
{
    PPROTECTED_DATA protected_data = NULL;
    return;

    protected_data = (PPROTECTED_DATA)ExAllocatePoolWithTag(
        NonPagedPool,
        DataSize + (ULONG)&(RTL_FIELD_TYPE(PROTECTED_DATA, mirror_data)),
        PROTECTED_DATA_TAG);

    RtlCopyMemory(protected_data->mirror_data, DataAddress, DataSize);
    protected_data->data_address = DataAddress;
    protected_data->data_size = DataSize;

    InsertTailList(&m_protected_data_list, &protected_data->list_entry);
}

VOID stop_protection()
{
    PPROTECTED_DATA protected_data = NULL;
    PLIST_ENTRY list_entry;

    return ;

    m_is_protected = FALSE;

    ZwWaitForSingleObject(m_thread_guard, FALSE, NULL);
    ZwWaitForSingleObject(m_thread_back_guard, FALSE, NULL);

    while ((list_entry = RemoveTailList(&m_protected_data_list))
            != &m_protected_data_list)
    {
        protected_data = CONTAINING_RECORD(
            list_entry, 
            PROTECTED_DATA, 
            list_entry);

        ExFreePoolWithTag(protected_data, PROTECTED_DATA_TAG);
    }
}

VOID _guard_thread( 
    IN PVOID StartContext 
    )
{
    NTSTATUS status = STATUS_SUCCESS;
    OBJECT_ATTRIBUTES thread_attribute = {0};
    KEVENT wait_event;
    LARGE_INTEGER wait_time;
    PLIST_ENTRY list_entry;
    PPROTECTED_DATA protected_data;
    ULONG old_status = 0;

    wait_time.QuadPart = -10000*1000*5;

    InitializeObjectAttributes(
        &thread_attribute, 
        NULL, 
        OBJ_KERNEL_HANDLE, 
        NULL, 
        NULL);

    KeInitializeEvent(&wait_event, SynchronizationEvent, FALSE);

    KeWaitForSingleObject(&wait_event, Executive, KernelMode, FALSE, &wait_time);

    while (m_is_protected)
    {
        status = ZwWaitForSingleObject(
            m_thread_back_guard, 
            FALSE, 
            &wait_time);
        if (status != STATUS_TIMEOUT)
        {
            status = PsCreateSystemThread(
                &m_thread_back_guard,
                THREAD_ALL_ACCESS,
                &thread_attribute,
                NULL,
                NULL,
                _back_guard_thread,
                NULL);

            ++ m_meltdown_count;
        }

        for (list_entry = m_protected_data_list.Flink; 
            list_entry != &m_protected_data_list;
            list_entry = list_entry->Flink)
        {
            protected_data = CONTAINING_RECORD(
                list_entry, 
                PROTECTED_DATA, 
                list_entry);

            if (protected_data->data_size != RtlCompareMemory(
                    protected_data->data_address,
                    protected_data->mirror_data, 
                    protected_data->data_size))
            {
                __asm
                {
                    cli;
                    mov eax, cr0;
                    mov old_status, eax;
                    and eax, 0xfffeffff;
                    mov cr0, eax;
                }

                RtlCopyMemory(
                    protected_data->data_address,
                    protected_data->mirror_data,
                    protected_data->data_size);

                __asm
                {
                    mov eax, old_status;
                    mov cr0, eax;
                    sti;
                }
                ++ m_meltdown_count;
            }
        }

        if (m_meltdown_count >= 0x20)
        {
            KeBugCheckEx(0x70, 0xc0000005, 0, 0, 0);
        }
    }
}

VOID _back_guard_thread( 
    IN PVOID StartContext 
    )
{
    NTSTATUS status = STATUS_SUCCESS;
    OBJECT_ATTRIBUTES thread_attribute = {0};
    KEVENT wait_event;
    LARGE_INTEGER wait_time;

    wait_time.QuadPart = -10000*1000*5;

    InitializeObjectAttributes(
        &thread_attribute, 
        NULL, 
        OBJ_KERNEL_HANDLE, 
        NULL, 
        NULL);

    KeInitializeEvent(&wait_event, SynchronizationEvent, FALSE);

    KeWaitForSingleObject(&wait_event, Executive, KernelMode, FALSE, &wait_time);

    while (m_is_protected)
    {
        
        status = ZwWaitForSingleObject(
            m_thread_guard, 
            FALSE, 
            &wait_time);
        if (status != STATUS_TIMEOUT)
        {
            status = PsCreateSystemThread(
                &m_thread_guard,
                THREAD_ALL_ACCESS,
                &thread_attribute,
                NULL,
                NULL,
                _guard_thread,
                NULL);
            ++ m_meltdown_count;
        }

        if (m_meltdown_count >= 0x20)
        {
            KeBugCheckEx(0x70, 0xc0000005, 0, 0, 0);
        }
    }
}